{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7971e3a7-e80c-45bd-b39d-6b922e6f2c80",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-21T02:31:59.130616Z",
     "iopub.status.busy": "2025-11-21T02:31:59.130389Z",
     "iopub.status.idle": "2025-11-21T02:31:59.132800Z",
     "shell.execute_reply": "2025-11-21T02:31:59.132395Z",
     "shell.execute_reply.started": "2025-11-21T02:31:59.130600Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 轻量级数据库\n",
    "import sqlite3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "25cc6c8f-e104-4793-9ae6-f9c3cb3935d0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-21T02:32:02.289459Z",
     "iopub.status.busy": "2025-11-21T02:32:02.289253Z",
     "iopub.status.idle": "2025-11-21T02:32:02.307503Z",
     "shell.execute_reply": "2025-11-21T02:32:02.307057Z",
     "shell.execute_reply.started": "2025-11-21T02:32:02.289443Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<sqlite3.Cursor at 0x7f609b7d33c0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 打开链接\n",
    "conn = sqlite3.connect(\"test.db\")\n",
    "# 数据库操作的句柄\n",
    "cursor = conn.cursor()\n",
    "\n",
    "cursor.execute(\"\"\"\n",
    "CREATE TABLE IF NOT EXISTS  employees (\n",
    "    id INTEGER PRIMARY KEY,\n",
    "    name TEXT,\n",
    "    department TEXT,\n",
    "    salary INTEGER\n",
    ")\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b82ad603-6b6b-4429-a939-fd93f00a24b3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sample_data = [\n",
    "    (6, \"王五\", \"开发部\", 32000),\n",
    "    ( 7, \"张三\", \"销售部\", 20000),\n",
    "    ( 8, \"赵六\", \"开发部\", 33000),\n",
    "    ( 9, \"李四\", \"销售部\", 15000)\n",
    "]\n",
    "cursor.executemany('INSERT INTO employees VALUES (?, ?, ?, ?)', sample_data)\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cb0d6f40-f092-43ce-b3f2-a340f6d8b15d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-11-21T02:46:15.251162Z",
     "iopub.status.busy": "2025-11-21T02:46:15.250957Z",
     "iopub.status.idle": "2025-11-21T02:46:15.813177Z",
     "shell.execute_reply": "2025-11-21T02:46:15.812629Z",
     "shell.execute_reply.started": "2025-11-21T02:46:15.251146Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "client = OpenAI(\n",
    "    api_key='sk-c031736e926b4608ba8603be5750d572',\n",
    "    base_url='https://api.deepseek.com/v1'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0128fbe2-e055-496a-99ed-7d3dd609d385",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2025-11-21T02:50:27.003016Z",
     "iopub.status.busy": "2025-11-21T02:50:27.002806Z",
     "iopub.status.idle": "2025-11-21T02:50:27.007255Z",
     "shell.execute_reply": "2025-11-21T02:50:27.006783Z",
     "shell.execute_reply.started": "2025-11-21T02:50:27.003002Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "数据库Schema:\n",
      "CREATE TABLE EMPLOYEES (\n",
      "id INTEGER\n",
      "name TEXT\n",
      "department TEXT\n",
      "salary INTEGER\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# 通过table_info 拿到employees 表的描述\n",
    "# llm 生成sql 提供上下文\n",
    "schema = cursor.execute(\"PRAGMA table_info(employees)\").fetchall()\n",
    "schema_str = \"CREATE TABLE EMPLOYEES (\\n\" + \"\\n\".join([f\"{col[1]} {col[2]}\" for col in schema]) + \"\\n)\"\n",
    "print(\"数据库Schema:\")\n",
    "print(schema_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d27e2097-089a-4380-8633-3b6c29a4f997",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2025-11-21T03:07:34.726239Z",
     "iopub.status.busy": "2025-11-21T03:07:34.726001Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def ask_deepseek(query, schema):\n",
    "    prompt = f\"\"\"\n",
    "    这是一个数据库的Schema:\n",
    "    {schema}\n",
    "    根据这个Schema,你能输出一个SQL查询来回答以下问题吗？\n",
    "    只输出SQL查询,不要输出任何内容，也不要带任何格式\n",
    "    问题：{query}\n",
    "    \"\"\"\n",
    "    # print(prompt)\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"deepseek-reasoner\",\n",
    "        max_tokens=2048,\n",
    "        messages=[{\n",
    "            \"role\": \"user\",\n",
    "            \"content\": prompt\n",
    "        }]\n",
    "    )\n",
    "    return response.choices[0].message\n",
    "    \n",
    "# ask_deepseek(\"开发部部门员工的姓名和工资是多少？\", schema)\n",
    "question = \"开发部部门员工的姓名和工资是多少？\"\n",
    "sql = ask_deepseek(question, schema_str)\n",
    "print(\"生成的sql查询：\")\n",
    "print(sql)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
